"""
@Project    : crnn
@Module     : json2crop.py
@Author     : wangqinggang@haier.com
@Created    : 2020/11/12 9:34
@Desc       : 航空行程crnn标注，得到json文件，把标注的数据按照取外接矩形裁剪，分别等到一个文件夹和一个txtx文件
            : 文件夹存放裁剪的图片
            : txt包含了图片名称和标签
"""

import numpy as np
import cv2
from glob import glob
from tqdm import tqdm
import os
import json
import random

from itertools import chain

'''裁剪图片'''


def random_sacle(img, box, left, right):
    bbox = box.copy()
    x_left = random.uniform(left, right)
    x_right = random.uniform(left, right)
    y_up = random.uniform(left, right)
    y_down = random.uniform(left, right)
    # x1,y1
    bbox[0][0] = max(0, bbox[0][0] - x_left)
    bbox[0][1] = max(0, bbox[0][1] - y_up)
    # x2,y2
    bbox[1][0] = min(bbox[1][0] + x_right, img.shape[1] - 1)
    bbox[1][1] = max(0, bbox[1][1] - y_up)
    # x3,y3
    bbox[2][0] = min(bbox[2][0] + x_right, img.shape[1] - 1)
    bbox[2][1] = min(bbox[2][1] + y_down, img.shape[0] - 1)
    # x4,y4
    bbox[3][0] = max(0, bbox[3][0] - x_left)
    bbox[3][1] = min(bbox[3][1] + y_down, img.shape[0] - 1)

    return bbox


'''裁剪图片'''


def crop_img_box(img, box, count):
    if count > 0:
        # w = (int)(np.linalg.norm(box[0] - box[1]))
        # h = (int)(np.linalg.norm(box[0] - box[3]))
        # scale = min(w,h)
        # if scale/5>3:
        #     left = scale/5
        # else:
        #     left = 3
        # if scale/2>10:
        #     right = scale/2
        # else:
        #     right = 10
        box = random_sacle(img, box, -3, 7)
    w = (int)(np.linalg.norm(box[0] - box[1]))  # 求范数 sqrt(x1**2+x2**2)
    h = (int)(np.linalg.norm(box[0] - box[3]))
    width = w
    height = h
    # if h > w * 1.5:
    #     width = h
    #     height = w
    #     M = cv2.getPerspectiveTransform(np.float32(box),
    #                                     np.float32(
    #                                         np.array([[width, 0], [width, height], [0, height], [0, 0]])))  # 透视变换
    # else:
    M = cv2.getPerspectiveTransform(np.float32(box),
                                    np.float32(np.array([[0, 0], [width, 0], [width, height], [0, height]])))

    cropped_image = cv2.warpPerspective(img, M, (width, height))  # 透视变换函数
    return cropped_image


'''裁剪图片:最大值最小值'''


def crop_img_max_min(img, box):
    xmin = int(box[:, 0].min() - random.uniform(-3, 20))
    ymin = int(box[:, 1].min() - random.uniform(-3, 20))
    xmax = int(box[:, 0].max() + random.uniform(-3, 20))
    ymax = int(box[:, 1].max() + random.uniform(-3, 20))
    xmin, ymin, xmax, ymax = max(0, xmin), max(0, ymin), min(xmax, img.shape[1] - 1), min(ymax, img.shape[0] - 1)
    crop_img = img[ymin:ymax, xmin:xmax]
    return crop_img


# json转txt
def json2txt(path, f):
    file, ext = os.path.splitext(path)
    txtpath = file.replace('json', 'txt') + '.txt'
    imgname = file.replace('json', 'img') + '.jpg'
    img = cv2.imread(imgname)
    # with open(txtpath, 'w+', encoding='utf-8') as f:
    with open(path, 'r', encoding='utf-8') as json_file:
        result = json.load(json_file)
        shape = result['shapes']
        for i in range(len(shape)):
            sh = shape[i]
            # print(sh)
            point = np.array(sh['points'])
            for j in range(2):
                crop_name = file.replace('json', 'crop_img') + '_' + str(i) + '_' + str(j) + '.jpg'
                # if '00569944_20190428173050001_11_0' in crop_name:
                #     print(crop_name)
                nerimg = crop_img_max_min(img, point)
                cv2.imwrite(crop_name, nerimg)
                pointstr = crop_name.split('\\')[-1] + ' ' + str(sh['label']).replace('：', ':').replace(' ', '') + '\n'
                f.write(pointstr)


def json2crop_crnn(path, f1, save_fd, key='test'):
    file, ext = os.path.splitext(path)
    imgname = file.replace('json', 'img') + '.jpg'
    if not os.path.exists(imgname):
        return
    img = cv2.imread(imgname)
    with open(path, 'r', encoding='utf-8') as json_file:
        result = json.load(json_file)
        shape = result['shapes']
        for i in range(len(shape)):
            sh = shape[i]
            point = np.array(sh['points'])
            label = sh['label']
            if len(label) == 1:
                crop_count = 8
            elif len(label) == 2:
                crop_count = 8
            else:
                crop_count = 5
            if key == 'test':
                crop_count = 1
            for j in range(crop_count):
                nerimg = crop_img_box(img, point, j)
                save_dir = os.path.join(save_fd, key)
                if not os.path.exists(save_dir):
                    os.mkdir(save_dir)
                crop_name = os.path.join(save_dir, os.path.basename(file)) + '_' + str(i) + '_' + str(j) + '.jpg'
                cv2.imwrite(crop_name, nerimg)
                pointstr = crop_name.split('\\')[-1] + ' ' + str(label).replace('：', ':').replace(' ', '') + '\n'
                pointstr = pointstr.replace('·', '.')
                # if '·' in pointstr:
                #     print(pointstr)
                f1.write(pointstr)


class JsonToDataset(object):
    def __init__(self, root_fd, save_fd, ratio=0.9):
        self.root_fd, self.save_fd = root_fd, save_fd
        self.ratio = ratio
        self.jsons = glob(self.root_fd+'/*.json')

    def gen_data(self, datasets, data_key):
        with open(os.path.join(self.save_fd, data_key+'.txt'), 'w+', encoding='utf-8') as f:
            for json_path in tqdm(datasets):
                json2crop_crnn(json_path, f, self.save_fd, data_key)
            f.close()

    def sample(self):
        data_len = len(self.jsons)
        random.shuffle(self.jsons)
        train_datasets, test_datasets = self.jsons[data_len//10:], self.jsons[:data_len//10]
        self.gen_data(train_datasets, 'train')
        self.gen_data(test_datasets, 'test')


if __name__ == '__main__':
    # imgs = r'D:\datasets\haier\piaoju\nlp\crnn\czc-ocr-labeled'
    # jsons = r'D:\datasets\haier\piaoju\nlp\crnn\czc-ocr-labeled'
    # json_list = glob(jsons + '/*.json')
    # with open(r'D:\datasets\haier\piaoju\nlp\crnn\\' + 'all.txt', 'w+', encoding='utf-8') as f1:
    #     for json_path in tqdm(json_list):
    #         json2crop_crnn(json_path, f1)
    #     f1.close()

    dataset = JsonToDataset(r'D:\datasets\haier\piaoju\nlp\crnn\czc-ocr-labeled', r'D:\datasets\haier\piaoju\nlp\crnn\czc-crnn-crop')
    dataset.sample()
